/* vim: set expandtab tabstop=4 softtabstop=4 shiftwidth=4: */
/* Copyright (C) 2012 Zongyou Yao. All rights reserved.     */

#include <ctype.h>
#include <limits.h>
#include <stdlib.h>
#include <string.h>
#include "assert.h"
#include "ap.h"
#include "fmt.h"
#include "xp.h"
#include "mem.h"

struct ap_t {
    int sign;
    int ndigits;
    int size;
    xp_t digits;
};

#define iszero(x) ((x)->ndigits == 1 && (x)->digits[0] == 0)
#define maxdigits(x, y) ((x)->ndigits > (y)->ndigits ? \
            (x)->ndigits : (y)->ndigits)
#define isone(x) ((x)->ndigits == 1 && (x)->digits[0] == 1)

static ap_t normalize(ap_t z, int n);
static int cmp(ap_t x, ap_t y);
static ap_t mk(int size)
{
    ap_t z = CALLOC(1, sizeof(*z) + size);
    assert(size > 0);
    z->sign = 1;
    z->size = size;
    z->ndigits = 1;
    z->digits = (xp_t)(z + 1);
    return z;
}

static ap_t set(ap_t z, long int n)
{
    if (n == LONG_MIN) {
        xp_fromint(z->size, z->digits, LONG_MAX + 1UL);
    } else if (n < 0) {
        xp_fromint(z->size, z->digits, -n);
    } else {
        xp_fromint(z->size, z->digits, n);
    }
    z->sign = n < 0 ? -1 : 1;
    return normalize(z, z->size);
}

static ap_t normalize(ap_t z, int n)
{
    z->ndigits = xp_length(n, z->digits);
    return z;
}

static ap_t add(ap_t z, ap_t x, ap_t y)
{
    int n = y->ndigits;
    if (x->ndigits < n) {
        return add(z, y, x);
    } else if (x->ndigits > n) {
        int carry = xp_add(n, z->digits, x->digits, y->digits, 0);
        z->digits[z->size - 1] = xp_sum(x->ndigits - n, &z->digits[n], &x->digits[n], carry);
    } else {
        z->digits[n] = xp_add(n, z->digits, z->digits, y->digits, 0);
    }

    return normalize(z, z->size);
}

static ap_t sub(ap_t z, ap_t x, ap_t y)
{
    int borrow, n = y ->ndigits;
    borrow = xp_sub(n, z->digits, x->digits, y->digits, 0);
    if (x->ndigits > n) {
        borrow = xp_diff(x->ndigits - n, &z->digits[n], &x->digits[n], borrow);
    }
    assert(borrow == 0);
    return normalize(z, z->size);
}

static ap_t mulmod(ap_t x, ap_t y, ap_t p)
{
    ap_t z, xy = ap_mul(x, y);
    z = ap_mod(xy, p);
    ap_free(&xy);
    return z;
}

static int cmp(ap_t x, ap_t y)
{
    if (x->ndigits != y->ndigits) {
        return x->ndigits - y->ndigits;
    } else {
        return xp_cmp(x->ndigits, x->digits, y->digits);
    }
}

ap_t ap_new(long int n)
{
    return set(mk(sizeof(long int)), n);
}

void ap_free(ap_t *z)
{
    assert(z && *z);
    FREE(*z);
}

ap_t ap_neg(ap_t x)
{
    ap_t z;
    assert(x);
    z = mk(x->ndigits);
    memcpy(z->digits, x->digits, x->ndigits);
    z->ndigits = x->ndigits;
    z->size = iszero(z) ? 1 : -x->sign;
    return z;
}

ap_t ap_mul(ap_t x, ap_t y)
{
    ap_t z;
    assert(x);
    assert(y);
    z = mk(x->ndigits + y->ndigits);
    xp_mul(z->digits, x->ndigits, x->digits, y->ndigits, y->digits);
    normalize(z, z->size);
    z->size = iszero(z) || ((x->size ^ y->size) == 0) ? 1 : -1;
    return z;
}

ap_t ap_add(ap_t x, ap_t y)
{
    ap_t z;
    assert(x);
    assert(y);
    if (((x->sign ^ y->sign) == 0)) {
        z = add(mk(maxdigits(x, y) + 1), x, y);
        z->sign = iszero(z) ? 1 : x->sign;
    } else {
        if (cmp(x, y) > 0) {
            z = sub(mk(x->ndigits), x, y);
            z->sign = iszero(z) ? 1 : x->sign;
        } else {
            z = sub(mk(y->ndigits), y, x);
            z->sign = iszero(z) ? 1 : -x->sign;
        }
    }
    return z;
}

ap_t ap_sub(ap_t x, ap_t y)
{
    ap_t z;
    assert(x);
    assert(y);
    if (!((x->sign ^ y->sign) == 0)) {
        z = add(mk(maxdigits(x, y) + 1), x, y);
        z->sign = iszero(z) ? 1 : x->sign;
    } else {
        if (cmp(x, y) > 0) {
            z = sub(mk(x->ndigits), x, y);
            z->sign = iszero(z) ? 1 : x->sign;
        } else {
            z = sub(mk(y->ndigits), y, x);
            z->sign = iszero(z) ? 1 : -x->sign;
        }
    }
    return z;
}

ap_t ap_div(ap_t x, ap_t y)
{
    ap_t q, r;
    assert(x);
    assert(y);
    assert(!iszero(y));
    q = mk(x->ndigits);
    r = mk(y->ndigits);
    {
        xp_t tmp = ALLOC(x->ndigits + y->ndigits + 2);
        xp_div(x->ndigits, q->digits, x->digits, y->ndigits, y->digits, r->digits, tmp);
        FREE(tmp);
    }
    normalize(q, q->size);
    normalize(r, r->size);
    q->sign = iszero(q) || ((x->sign ^ y->sign) == 0) ? 1 : -1;
    if (!((x->sign ^ y->sign) == 0) && !iszero(r)) {
        int carry = xp_sum(q->size, q->digits, q->digits, 1);
        assert(carry == 0);
        normalize(q, q->size);
    }
    ap_free(&r);
    return q;
}

ap_t ap_mod(ap_t x, ap_t y)
{
    ap_t q, r;
    assert(x);
    assert(y);
    assert(!iszero(y));
    q = mk(x->ndigits);
    r = mk(y->ndigits);
    {
        xp_t tmp = ALLOC(x->ndigits + y->ndigits + 2);
        xp_div(x->ndigits, q->digits, x->digits, y->ndigits, y->digits, r->digits, tmp);
        FREE(tmp);
    }
    normalize(q, q->size);
    normalize(r, r->size);
    q->size = iszero(q) || ((x->sign ^ y->sign) == 0) ? 1 : -1;
    if (!((x->sign ^ y->sign) == 0) && !iszero(r)) {
        int borrow = xp_sub(r->size, r->digits, y->digits, r->digits, 0);
        assert(borrow == 0);
        normalize(r, r->size);
    }
    ap_free(&q);
    return r;
}

ap_t ap_pow(ap_t x, ap_t y, ap_t p)
{
    ap_t z;
    assert(x);
    assert(y);
    assert(y->sign == 1);
    assert(!p || (p->sign == 1 && !iszero(p) && !isone(p)));
    if (iszero(x)) {
        return ap_new(0);
    }
    if (iszero(y)) {
        return ap_new(1);
    }
    if (isone(x)) {
        return ap_new((((y)->digits[0] & 1) == 0) ? 1 : x->sign);
    }
    if (p) {
        if (isone(y)) {
            z = ap_mod(x, p);
        } else {
            ap_t y2 = ap_rshift(y, 1), t = ap_pow(x, y2, p);
            z = mulmod(t, t, p);
            ap_free(&y2);
            ap_free(&t);
            if (!(((y)->digits[0] & 1) == 0)) {
                z = mulmod(y2 = ap_mod(x, p), t = z, p);
                ap_free(&y2);
                ap_free(&t);
            }
        }
    } else {
        if (isone(y)) {
            z = ap_addi(x, 0);
        } else {
            ap_t y2 = ap_rshift(y, 1), t = ap_pow(x, y2, NULL);
            z = ap_mul(t, t);
            ap_free(&y2);
            ap_free(&t);
            if (!(((y)->digits[0] & 1) == 0)) {
                z = ap_mul(x, t = z);
                ap_free(&t);
            }
        }
    }
    return z;
}

int ap_cmp(ap_t x, ap_t y)
{
    assert(x);
    assert(y);
    if (!((x->sign ^ y->sign) == 0)) {
        return x->sign;
    } else if (x->sign == 1) {
        return cmp(x, y);
    } else {
        return cmp(y, x);
    }
}

ap_t ap_addi(ap_t x, long int y)
{
    unsigned char d[sizeof(unsigned long)];
    struct ap_t t;
    t.size = sizeof(d);
    t.digits = d;
    return ap_add(x, set(&t, y));
}

ap_t ap_subi(ap_t x, long int y)
{
    unsigned char d[sizeof(unsigned long)];
    struct ap_t t;
    t.size = sizeof(d);
    t.digits = d;
    return ap_sub(x, set(&t, y));
}

ap_t ap_muli(ap_t x, long int y)
{
    unsigned char d[sizeof(unsigned long)];
    struct ap_t t;
    t.size = sizeof(d);
    t.digits = d;
    return ap_mul(x, set(&t, y));
}

ap_t ap_divi(ap_t x, long int y)
{
    unsigned char d[sizeof(unsigned long)];
    struct ap_t t;
    t.size = sizeof(d);
    t.digits = d;
    return ap_div(x, set(&t, y));
}

int ap_cmpi(ap_t x, long int y)
{
    unsigned char d[sizeof(unsigned long)];
    struct ap_t t;
    t.size = sizeof(d);
    t.digits = d;
    return ap_cmp(x, set(&t, y));
}

long int ap_modi(ap_t x, long int y)
{
    long int rem;
    ap_t r;
    unsigned char d[sizeof(unsigned long)];
    struct ap_t t;
    t.size = sizeof(d);
    t.digits = d;
    r = ap_mod(x, set(&t, y));
    rem = xp_toint(r->ndigits, r->digits);
    ap_free(&r);
    return rem;
}

ap_t ap_lshift(ap_t x, int s)
{
    ap_t z;
    assert(x);
    assert(s >= 0);
    z = mk(x->ndigits + ((s + 7) & ~7) / 8);
    xp_lshift(z->size, z->digits, x->ndigits, x->digits, s, 0);
    z->sign = x->sign;
    return normalize(z, z->size);
}

ap_t ap_rshift(ap_t x, int s)
{
    assert(x);
    assert(s >= 0);
    if (s >= 8 * x->ndigits) {
        return ap_new(0);
    } else {
        ap_t z = mk(x->ndigits - s / 8);
        xp_rshift(z->size, z->digits, x->ndigits, x->digits, s, 0);
        normalize(z, z->size);
        z->sign = iszero(z) ? 1 : x->sign;
        return z;
    }
}

long int ap_toint(ap_t x)
{
    unsigned long u;
    assert(x);
    u = xp_toint(x->ndigits, x->digits) % (LONG_MAX + 1UL);
    if (x->sign == -1) {
        return -(long)u;
    } else {
        return (long)u;
    }
}

ap_t ap_fromstr(const char *str, int base, char **end)
{
    ap_t z;
    const char *p = str;
    char *endp, sign = '\0';
    int carry;
    assert(p);
    assert(base >= 2 && base <= 36);
    while (*p && isspace(*p)) {
        p++;
    }
    if (*p == '-' || *p == '+') {
        sign = *p++;
    }
    {
        const char *start;
        int k, n = 0;
        for ( ; *p == '0' && p[1] == '0'; p++) {
            ;
        }
        start = p;
        for ( ; (('0' <= *p && *p <= '9' && *p < '0' + base) ||
                 ('a' <= *p && *p <= 'z' && *p < 'a' + base - 10) ||
                 ('A' <= *p && *p <= 'Z' && *p < 'A' + base - 10)); ++p) {
            n++;
        }
        for (k = 1; (1 << k) < base; ++k) {
            ;
        }
        z = mk(((k * n + 7) & ~7) / 8);
        p = start;
    }
    carry = xp_fromstr(z->size, z->digits, p, base, &endp);
    assert(carry == 0);
    normalize(z, z->size);
    if (endp == p) {
        endp = (char *)str;
        z = ap_new(0);
    } else {
        z->sign = iszero(z) || sign != '-' ? 1 : -1;
    }
    if (end) {
        *end = (char *)endp;
    }
    return z;
}

char *ap_tostr(char *str, int size, int base, ap_t x)
{
    xp_t q;
    assert(x);
    assert(base >= 2 && base <= 36);
    assert(str == NULL || size > 1);
    if (str == NULL) {
        {
        int k;
        for (k = 5; (1 << k) > base; k--) {
            ;
        }
        size = (8 * x->ndigits) / k + 1 + 1;
        if (x->sign == -1) {
            size++;
        }
        }
        str = ALLOC(size);
    }
    q = ALLOC(x->ndigits);
    memcpy(q, x->digits, x->ndigits);
    if (x->sign == -1) {
        str[0] = '-';
        xp_tostr(str + 1, size - 1, base, x->ndigits, q);
    } else {
        xp_tostr(str, size, base, x->ndigits, q);
    }
    FREE(q);
    return str;
}

void ap_fmt(int code, va_list_box *box, int put(int c, void *cl), void *cl,
            unsigned char flags[], int width, int precision)
{
    ap_t x;
    char *buf;
    assert(box && flags);
    x = va_arg(box->ap, ap_t);
    assert(x);
    buf = ap_tostr(NULL, 0, 10, x);
    fmt_putd(buf, strlen(buf), put, cl, flags, width, precision);
    FREE(buf);
}
